// ++++++++++++++++++++++++++++++++++++++++
// 《Go 语言+Vue.js 实战派——基于 Gin 框架》源码
// ++++++++++++++++++++++++++++++++++++++++
// 作者公众号：源码大数据
// 自媒体账号（抖音、视频号、快手、B站、知乎）：廖显东-ShirDon
// 仓库地址：https://gitee.com/shirdonl/goGinVue.git
// 仓库地址：https://github.com/shirdonl/goGinVue.git
// 博客：https://www.shirdon.com
// 购买图书：https: //item.jd.com/14958064.html
// ++++++++++++++++++++++++++++++++++++++++

package main

import (
	"github.com/gin-gonic/gin"
	"net/http"
	"sync"
	"time"
)

type RateLimiter struct {
	sync.Mutex
	Requests map[string]int
}

func main() {
	r := gin.Default()

	//设置速率限制中间件
	rateLimiter := &RateLimiter{Requests: make(map[string]int)}
	r.Use(RateLimitMiddleware(rateLimiter, 2, 5*time.Minute))

	//定义路由
	r.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "Hi, welcome!")
	})

	r.Run(":8080")
}

func RateLimitMiddleware(rateLimiter *RateLimiter, limit int, interval time.Duration) gin.HandlerFunc {
	return func(c *gin.Context) {
		rateLimiter.Lock()
		defer rateLimiter.Unlock()

		ip := c.ClientIP()
		requests := rateLimiter.Requests[ip]

		if requests > limit {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "Too many requests."})
			return
		}

		rateLimiter.Requests[ip] = requests + 1
		time.AfterFunc(interval, func() {
			rateLimiter.Lock()
			defer rateLimiter.Unlock()
			rateLimiter.Requests[ip] = 0
		})

		c.Next()
	}
}
